import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import os
from torchvision import transforms

"""
自己用的时候必须继承Dataset并改写其中的 __getitem__函数 和 __len__函数
"""


class torch_Dataset_class(Dataset):
    """
    介绍关于Dataset类的使用
    torch.utils.data.Dataset是一个抽象类，用户想要加载自定义的数据只需要继承这个类，并且覆写其中的两个方法即可
    1. __len__: 实现len(dataset)返回整个数据集的大小。
    2. __getitem__: 用来获取一些索引的数据，使dataset[i]返回数据集中第i个样本。
    不覆写这两个方法会直接返回错误。
    """

    def __init__(self):
        # 调用父类的构造函数，不过这不是必须的
        super().__init__()
        self.images = None
        self.path = None

    def __len__(self):
        """
        获取数据集中数据量大小的方法（必须要被重写）
        """
        return len(self.images)

    def __getitem__(self, item):
        """
        获取第i个数据的属性的方法（必须要被重写）
        """
        return self.images[item]

    def read_frompath(self, path):
        """
        输入数据集路径，载入名称，该函数返回的是名称，而不是图片
        """
        self.path = path
        self.images = os.listdir(path)
